from openai import OpenAI
from typing import Literal
# Set OpenAI's API key and API base to use vLLM's API server.
openai_api_key = "EMPTY"
openai_api_base = "http://127.0.0.1:8007/v1"

class qwen2ForOpenAIv2(object):
    """
    使用OpenAI API接口实现的分类服务（tochat/resource）
    """
    def __init__(self, template_path: str, model_name: str = "xiaoxingv3"):
        # OpenAI客户端配置
        self.client = OpenAI(
            api_key=openai_api_key,
            base_url=openai_api_base
        )
        self.model_name = model_name
        self.template_path = template_path
        self.INSTRUCTION = self._load_template(template_path)
        
        # 映射原始生成参数到OpenAI参数
        self.generation_config = {
            "temperature": 0.1,      # 对应原始do_sample=False
            "max_tokens": 5,         # 对应max_new_tokens=5
            # 重复惩罚参数需要转换公式
            "frequency_penalty": 0.18  # 经验值，近似repetition_penalty=1.1
        }

    def _load_template(self, template_path):
        # 保持与原实现一致
        with open(template_path, 'r', encoding='utf-8') as f:
            INSTRUCTION = f.read().strip()
            print(INSTRUCTION)
        return INSTRUCTION

    def build_chat_messages(self, user_input):
        """构建OpenAI兼容的对话格式"""
        return [
            {
                "role": "system",
                "content": "你是一位精通自然语言理解的专家，专注于精准分析用户输入并识别其意图，擅长处理中文语境中的细微语义差异。"
            },
            {
                "role": "user",
                "content": f"{self.INSTRUCTION}\n{user_input}"
            }
        ]

    def classify(self, text) -> Literal["resource", "tochat"]:
        messages = self.build_chat_messages(text)
        
        try:
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=messages, # type: ignore
                **self.generation_config
            ) # type: ignore
            # 提取生成的文本内容
            response_text = response.choices[0].message.content
        except Exception as e:
            print(f"API调用失败: {str(e)}")
            response_text = ""  # 返回空字符串触发后处理逻辑
        
        return self.postprocess(response_text)

    def postprocess(self, text):
        """保持与原实现一致的后处理逻辑"""
        clean = text.strip().lower().replace("`", "").replace(" ", "").replace("\n", "")
        if clean == "resource":
            return "resource"
        elif clean == "tochat":
            return "tochat"
        else:
            return "tochat"

